import cv2
import sys
import numpy as np
import datetime
import os
import glob
from retinaface import RetinaFace
import time
os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT'] = '0'

#If you find the speed of face detection is slow, you can type "set MXNET_CUDNN_AUTOTUNE_DEFAULT=0" in your terminal.

def retinaface_detect(img, detector):
	""" Face detection based on RetinaFace model
		Parameters:
		img: img for face detection
		detector: detector of retinaface model
		return: location of faces and landmarks
	"""
	start_time = time.time()

	thresh = 0.8
	scales = [1024, 1980]
	#scales = [512, 990]
	count = 1
	#gpuid = 0

	#detector = RetinaFace('./model/R50', 0, gpuid, 'net3')
	#detector = RetinaFace('./models/mnet.25/mnet.25', 0, gpuid, 'net3')
	#detector = RetinaFace('./models/retinaface-R50-cudnnoff/R50', 0, gpuid, 'net3')

	# count time of loading model
	loadmodel_time = time.time()
	#print('1. Load model time: %s s' % (loadmodel_time - start_time))

	#print(img.shape)
	im_shape = img.shape
	target_size = scales[0]
	max_size = scales[1]
	im_size_min = np.min(im_shape[0:2])
	im_size_max = np.max(im_shape[0:2])

	#im_scale = 1.0
	#if im_size_min>target_size or im_size_max>max_size:
	im_scale = float(target_size) / float(im_size_min)
	# prevent bigger axis from being more than max_size:
	if np.round(im_scale * im_size_max) > max_size:
		im_scale = float(max_size) / float(im_size_max)

	print('im_scale', im_scale)
	scales = [im_scale]
	flip = False

	# count setting all parameters time
	setting_time = time.time()
	print('2. Setting Parameters time: %s s' % (setting_time - loadmodel_time))

	for c in range(count):
		faces, landmarks = detector.detect(img, thresh, scales=scales, do_flip=flip)
		print(c, faces.shape, landmarks.shape)

	# count face detection time
	facedetect_time = time.time()
	print('3. Face detection time: %s s' % (facedetect_time - setting_time))

	return faces, landmarks

if __name__ == '__main__':
	s_time = time.time()

	# load the image
	img_path = 'imgs/t2.jpg'
	img = cv2.imread(img_path)
	#h = int(img.shape[0])
	#w = int(img.shape[1])
	#img = cv2.resize(img,(w,h))

	#print('image size: ', h, w)

	# load retina face model
	gpuid = 0
	detector = RetinaFace('./models/retinaface-R50-cudnnoff/R50', 0, gpuid, 'net3')

	# count load model time
	loadmodel_time = time.time()
	print('1. Load model time: %s s' % (loadmodel_time - s_time))
	
	# detect faces in the image
	faces, landmarks = retinaface_detect(img, detector)

	start0_time = time.time()

	# show the faces on the image
	if faces is not None:
		print('find', faces.shape[0], 'faces')
		for i in range(faces.shape[0]):
			#print('score', faces[i][4])
			box = faces[i].astype(np.int)
			#color = (255,0,0)
			color = (0,0,255)
			cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, 2)
			if landmarks is not None:
				landmark5 = landmarks[i].astype(np.int)
				#print(landmark.shape)
				for l in range(landmark5.shape[0]):
					color = (0,0,255)
					if l==0 or l==3:
						color = (0,255,0)
					cv2.circle(img, (landmark5[l][0], landmark5[l][1]), 1, color, 2)
	# count draw image time
	draw_time = time.time()
	print('4. Draw image time: %s s' % (draw_time - start0_time))

	filename = './results/detector_test.jpg'
	print('writing', filename)
	cv2.imwrite(filename, img)

	# count save image time
	save_time = time.time()
	print('5. Save image time: %s s' % (save_time - draw_time))

	# count the whole processing time
	end_time = time.time()
	print('6. The whole time: %s s' % (end_time - s_time))

	cv2.imshow('ret', img)

	cv2.waitKey(0)
	cv2.destroyAllWindows()